# -*- coding: utf-8 -*-

from ezds.ezdlearn.config import load_config
import os
from ezds.ezdlearn.utils.load import *
from torch.utils.data import DataLoader
import warnings
warnings.filterwarnings('ignore')

#%% main
if __name__ == "__main__":
    config_path = "configs/DINet_frame.yaml"
    if config_path is not None:
        cfg = load_config(config_path)
    else:
        cfg = load_config()
    basename = os.path.basename(cfg.trainer_py_path)
    Trainer = get_trainer_by_file(cfg.trainer_py_path)
    Dataset, input_format = get_dataset_by_file(cfg.dset_py_path)
    if cfg.model_name is None:
        cfg.model_name = basename.split('.')[0]
    # dataset phase 
    dataset = Dataset(cfg)
    dataloader = DataLoader(dataset, batch_size=cfg.batch_size, shuffle=cfg.is_shuffle, num_workers=cfg.n_workers)    
    if cfg.val_step > 0:
        val_dataset = Dataset(cfg, is_train=False)
        val_dataloader = DataLoader(val_dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=cfg.n_workers)
    else:
        val_dataloader = None
    # trainer phase
    trainer = Trainer(cfg)
    trainer.init_params()
    trainer.fit(dataloader, input_format, ignore_err=False, auto_load=True, use_jupyter=False, val_dataloader=val_dataloader)